import { BaseEmbeddings } from '@cherrystudio/embedjs-interfaces'

export interface SemanticSplitterConfig {
  // 最小分片大小（字符数）
  minChunkSize?: number
  // 最大分片大小（字符数）
  maxChunkSize?: number
  // 语义相似度阈值 (0-1)
  similarityThreshold?: number
  // 分片重叠长度
  overlap?: number
}

export class SemanticSplitter {
  private readonly embeddingModel: BaseEmbeddings
  private readonly config: Required<SemanticSplitterConfig>

  constructor(embeddingModel: BaseEmbeddings, config: SemanticSplitterConfig = {}) {
    this.embeddingModel = embeddingModel
    // 设置默认配置
    this.config = {
      minChunkSize: 100,
      maxChunkSize: 1000,
      similarityThreshold: 0.7,
      overlap: 20,
      ...config
    }
  }

  /**
   * 计算两个向量的余弦相似度
   */
  private cosineSimilarity(vec1: number[], vec2: number[]): number {
    if (vec1.length !== vec2.length) {
      throw new Error('向量维度不匹配')
    }

    let dotProduct = 0
    let norm1 = 0
    let norm2 = 0

    for (let i = 0; i < vec1.length; i++) {
      dotProduct += vec1[i] * vec2[i]
      norm1 += vec1[i] * vec1[i]
      norm2 += vec2[i] * vec2[i]
    }

    return dotProduct / (Math.sqrt(norm1) * Math.sqrt(norm2))
  }

  /**
   * 按标点符号分割文本为初始句子
   */
  private splitIntoSentences(text: string): string[] {
    // 移除多余空格
    text = text.replace(/ /g, '') // 移除此行，保留文本中的空格
    // 先按标点符号分割
    const rawSentences = text.split(/(?<=[。！？!?])\s*/).map((s) => s.trim())

    // 合并短句（长度小于10的句子）
    const mergedSentences: string[] = []
    for (let i = 0; i < rawSentences.length; i++) {
      const current = rawSentences[i]

      // 如果当前句子长度小于10且不是最后一个句子，则与下一个句子合并
      if (current.length < 10 && i < rawSentences.length - 1) {
        rawSentences[i + 1] = current + rawSentences[i + 1]
      } else {
        // 否则添加到结果中（长度大于等于10或是最后一个句子）
        mergedSentences.push(current)
      }
    }

    return mergedSentences
  }

  /**
   * 合并相邻的相似句子（优化：批量嵌入+缓存+并发）
   */
  private async mergeSimilarSentences(sentences: string[]): Promise<string[]> {
    if (sentences.length <= 1) return sentences

    // 1. 批量获取所有句子的 embedding，缓存结果
    const embeddingCache: number[][] = await this.getEmbeddingsBatch(sentences)

    const chunks: string[] = []
    let currentChunk = sentences[0]
    let currentLength = currentChunk.length
    let currentVec = embeddingCache[0]

    for (let i = 1; i < sentences.length; i++) {
      const currentSentence = sentences[i]
      const nextLength = currentLength + currentSentence.length
      const nextVec = embeddingCache[i]

      // 如果合并后超过最大长度，保存当前分片并开始新分片
      if (nextLength > this.config.maxChunkSize) {
        chunks.push(currentChunk)
        currentChunk = currentSentence
        currentLength = currentSentence.length
        currentVec = nextVec
        continue
      }

      // 计算相邻句子的语义相似度
      const similarity = this.cosineSimilarity(currentVec, nextVec)

      // 如果相似度高于阈值且长度合适，或者当前分片长度小于最小分片大小，则合并句子
      if (similarity >= this.config.similarityThreshold || currentLength < this.config.minChunkSize) {
        currentChunk += ' ' + currentSentence
        currentLength = nextLength
        // 合并后更新 currentVec（简单处理：用下一个句子的 embedding 作为当前 embedding）
        currentVec = nextVec
      } else {
        // 相似度低，保存当前分片并开始新分片
        chunks.push(currentChunk)
        currentChunk = currentSentence
        currentLength = currentSentence.length
        currentVec = nextVec
      }
    }

    // 添加最后一个分片
    if (currentChunk) {
      chunks.push(currentChunk)
    }

    return chunks
  }

  /**
   * 批量获取句子的 embedding，若模型不支持批量则降级为单条
   */
  private async getEmbeddingsBatch(sentences: string[]): Promise<number[][]> {
    if (typeof (this.embeddingModel as any).embedDocuments === 'function') {
      // 优先使用批量接口
      return await (this.embeddingModel as any).embedDocuments(sentences)
    } else {
      // 否则逐条获取
      return await Promise.all(sentences.map((s) => this.embeddingModel.embedQuery(s)))
    }
  }

  /**
   * 执行语义分片
   * @param text 输入文本
   * @returns 分片结果数组
   */
  async splitText(text: string): Promise<string[]> {
    // 1. 初始分句
    // 移除或注释掉日志
    // console.log('text:', text)
    const sentences = this.splitIntoSentences(text)
    // console.log('sentence:', sentences)

    // 2. 合并相似句子
    const mergedChunks = await this.mergeSimilarSentences(sentences)

    // 3. 处理分片重叠
    const finalChunks: string[] = []
    for (let i = 0; i < mergedChunks.length; i++) {
      const chunk = mergedChunks[i]
      if (i > 0 && this.config.overlap > 0) {
        // 添加前一个分片的结尾作为重叠
        const prevChunk = mergedChunks[i - 1]
        const overlap = prevChunk.slice(-this.config.overlap)
        finalChunks.push(overlap + chunk)
      } else {
        finalChunks.push(chunk)
      }
    }

    return finalChunks
  }
}
